#!/bin/bash

# =============================================
# SSH访问监控与自动封禁脚本 v1.0
# =============================================
# 功能特点：
# ✓ 监控SSH暴力破解攻击
# ✓ 检测异常登录尝试
# ✓ 自动封禁恶意IP地址
# ✓ 白名单保护机制
# ✓ 灵活的配置选项
# ✓ 定时自动解封
# ✓ 详细日志记录
# =============================================

# 检查是否为root用户
if [ "$(id -u)" -ne 0 ]; then
    echo "❌ 此脚本需要root权限运行"
    echo "请使用: sudo $0"
    exit 1
fi

# 定义颜色
RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
BLUE="\033[34m"
CYAN="\033[36m"
MAGENTA="\033[35m"
RESET="\033[0m"

# 配置文件路径
SCRIPT_DIR="/opt/ssh-monitor"
CONFIG_FILE="$SCRIPT_DIR/ssh-monitor.conf"
WHITELIST_FILE="$SCRIPT_DIR/whitelist.txt"
BANLIST_FILE="$SCRIPT_DIR/banlist.txt"
LOG_FILE="$SCRIPT_DIR/ssh-monitor.log"

# 创建必要的目录
mkdir -p "$SCRIPT_DIR"

# 默认配置
DEFAULT_CONFIG='# SSH监控配置文件

# 检测阈值 - 一小时内失败次数超过此值将被封禁
MAX_FAILED_ATTEMPTS=5

# 检测时间范围 - 统计多长时间内的失败次数（小时）
CHECK_PERIOD_HOURS=1

# 封禁时间（分钟）- 0表示永久封禁
BAN_DURATION=120

# SSH日志文件路径
SSH_LOG=""

# 监控的用户名 - 空表示监控所有用户
MONITOR_USERS=""

# 排除的IP段（CIDR格式，逗号分隔）
EXCLUDE_NETWORKS="127.0.0.0/8,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16"

# 是否启用邮件通知
ENABLE_EMAIL=false
EMAIL_TO="admin@example.com"

# 是否启用钉钉通知
ENABLE_DINGTALK=false
DINGTALK_WEBHOOK=""

# 检测敏感用户登录尝试
SENSITIVE_USERS="root,admin,administrator,mysql,postgres,oracle"
'

# 检测并初始化防火墙工具
detect_firewall() {
    local firewall_tool=""
    
    # 检测可用的防火墙工具
    if command -v iptables >/dev/null 2>&1; then
        firewall_tool="iptables"
        log_message "INFO" "检测到防火墙工具: iptables"
    elif command -v nft >/dev/null 2>&1; then
        firewall_tool="nftables"
        log_message "INFO" "检测到防火墙工具: nftables"
    elif command -v ufw >/dev/null 2>&1; then
        firewall_tool="ufw"
        log_message "INFO" "检测到防火墙工具: ufw"
    else
        log_message "ERROR" "未找到支持的防火墙工具"
        printf "${RED}❌ 系统中未找到防火墙管理工具${RESET}\n"
        printf "${YELLOW}Deepin系统建议安装iptables:${RESET}\n"
        printf "  ${CYAN}sudo apt update${RESET}\n"
        printf "  ${CYAN}sudo apt install iptables-persistent${RESET}\n\n"
        printf "${YELLOW}或者使用ufw (推荐):${RESET}\n"
        printf "  ${CYAN}sudo apt install ufw${RESET}\n"
        printf "  ${CYAN}sudo ufw enable${RESET}\n\n"
        
        printf "${MAGENTA}是否现在安装iptables? (y/n):${RESET} "
        read -r choice
        case "$choice" in
            [Yy]*)
                log_message "INFO" "正在安装iptables..."
                apt update && apt install -y iptables iptables-persistent
                if [ $? -eq 0 ]; then
                    firewall_tool="iptables"
                    log_message "INFO" "iptables安装成功"
                else
                    log_message "ERROR" "iptables安装失败"
                    return 1
                fi
                ;;
            *)
                log_message "WARN" "用户取消安装，脚本功能受限"
                return 1
                ;;
        esac
    fi
    
    # 设置全局防火墙工具变量
    FIREWALL_TOOL="$firewall_tool"
    return 0
}

# 检查IP是否已被封禁（兼容多种防火墙）
is_banned() {
    local ip="$1"
    
    case "$FIREWALL_TOOL" in
        "iptables")
            iptables -L INPUT -n | grep -q "$ip"
            ;;
        "nftables")
            nft list ruleset | grep -q "$ip"
            ;;
        "ufw")
            ufw status | grep -q "$ip"
            ;;
        *)
            log_message "ERROR" "不支持的防火墙工具: $FIREWALL_TOOL"
            return 1
            ;;
    esac
}

# 封禁IP（兼容多种防火墙）
ban_ip() {
    local ip="$1"
    local reason="$2"
    local failed_count="$3"
    
    # 检查IP格式
    if ! echo "$ip" | grep -E '^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$' >/dev/null; then
        log_message "ERROR" "无效的IP地址格式: $ip"
        return 1
    fi
    
    # 检查是否在白名单中
    if is_whitelisted "$ip"; then
        log_message "INFO" "IP $ip 在白名单中，跳过封禁"
        return 0
    fi
    
    # 检查是否已被封禁
    if is_banned "$ip"; then
        log_message "INFO" "IP $ip 已被封禁"
        return 0
    fi
    
    # 根据防火墙工具执行封禁
    local ban_success=0
    case "$FIREWALL_TOOL" in
        "iptables")
            iptables -I INPUT -s "$ip" -j DROP
            ban_success=$?
            ;;
        "nftables")
            # 创建nftables表和链（如果不存在）
            nft add table inet filter 2>/dev/null
            nft add chain inet filter input "{ type filter hook input priority 0; }" 2>/dev/null
            nft add rule inet filter input ip saddr "$ip" drop 2>/dev/null
            ban_success=$?
            ;;
        "ufw")
            ufw deny from "$ip" 2>/dev/null
            ban_success=$?
            ;;
        *)
            log_message "ERROR" "不支持的防火墙工具: $FIREWALL_TOOL"
            return 1
            ;;
    esac
    
    if [ $ban_success -eq 0 ]; then
        local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
        echo "$timestamp|$ip|$reason|失败次数:$failed_count|防火墙:$FIREWALL_TOOL" >> "$BANLIST_FILE"
        log_message "BAN" "已封禁IP: $ip, 原因: $reason (失败次数: $failed_count)"
        
        # 发送通知
        send_notification "🚫 IP封禁通知" "IP: $ip\n原因: $reason\n失败次数: $failed_count\n时间: $(date)"
        
        # 设置自动解封（如果配置了封禁时间）
        if [ "$BAN_DURATION" -gt 0 ]; then
            setup_auto_unban "$ip" "$BAN_DURATION"
        fi
        
        return 0
    else
        log_message "ERROR" "封禁IP失败: $ip"
        return 1
    fi
}

# 解封IP（兼容多种防火墙）
unban_ip() {
    local ip="$1"
    
    # 检查IP格式
    if ! echo "$ip" | grep -E '^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$' >/dev/null; then
        log_message "ERROR" "无效的IP地址格式: $ip"
        return 1
    fi
    
    # 根据防火墙工具执行解封
    local unban_success=0
    case "$FIREWALL_TOOL" in
        "iptables")
            # 删除所有相关的DROP规则
            while iptables -D INPUT -s "$ip" -j DROP 2>/dev/null; do
                continue
            done
            unban_success=0  # iptables -D 在没有规则时会返回非0，但这里我们认为是成功的
            ;;
        "nftables")
            # 删除nftables规则（这个比较复杂，需要先找到handle）
            local handles=$(nft -a list chain inet filter input | grep "$ip" | awk '{print $NF}')
            for handle in $handles; do
                nft delete rule inet filter input handle "$handle" 2>/dev/null
            done
            unban_success=0
            ;;
        "ufw")
            ufw delete deny from "$ip" 2>/dev/null
            unban_success=0  # ufw delete 在没有规则时也会报错，但我们认为是成功的
            ;;
        *)
            log_message "ERROR" "不支持的防火墙工具: $FIREWALL_TOOL"
            return 1
            ;;
    esac
    
    # 从封禁列表中移除
    if [ -f "$BANLIST_FILE" ]; then
        grep -v "|$ip|" "$BANLIST_FILE" > "$BANLIST_FILE.tmp" && mv "$BANLIST_FILE.tmp" "$BANLIST_FILE"
    fi
    
    log_message "INFO" "已解封IP: $ip"
    return 0
}

# 获取当前封禁的IP数量（兼容多种防火墙）
get_banned_count() {
    case "$FIREWALL_TOOL" in
        "iptables")
            iptables -L INPUT -n 2>/dev/null | grep -c "DROP.*[0-9]\+\.[0-9]\+\.[0-9]\+\.[0-9]\+" || echo "0"
            ;;
        "nftables")
            nft list ruleset 2>/dev/null | grep -c "ip saddr.*drop" || echo "0"
            ;;
        "ufw")
            ufw status numbered 2>/dev/null | grep -c "DENY.*[0-9]\+\.[0-9]\+\.[0-9]\+\.[0-9]\+" || echo "0"
            ;;
        *)
            echo "0"
            ;;
    esac
}

# 显示当前封禁的IP列表（兼容多种防火墙）
show_current_banned_ips() {
    case "$FIREWALL_TOOL" in
        "iptables")
            iptables -L INPUT -n 2>/dev/null | grep DROP | grep -E '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | awk '{print "  " $4}' | sed 's/\/32//'
            ;;
        "nftables")
            nft list ruleset 2>/dev/null | grep "ip saddr" | grep "drop" | grep -oE '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | awk '{print "  " $1}'
            ;;
        "ufw")
            ufw status numbered 2>/dev/null | grep DENY | grep -oE '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | awk '{print "  " $1}'
            ;;
        *)
            echo "  ${YELLOW}无法显示 - 防火墙工具不支持${RESET}"
            ;;
    esac
}

# 初始化配置文件
init_config() {
    if [ ! -f "$CONFIG_FILE" ]; then
        echo "$DEFAULT_CONFIG" > "$CONFIG_FILE"
        printf "${GREEN}✓ 已创建配置文件: $CONFIG_FILE${RESET}\n"
    fi
    
    # 初始化白名单
    if [ ! -f "$WHITELIST_FILE" ]; then
        cat > "$WHITELIST_FILE" << EOF
# SSH访问白名单IP地址，每行一个
# 支持单个IP和IP段（CIDR格式）
# 注释行以#开头

127.0.0.1
::1

# 示例：
# 192.168.1.100
# 10.0.0.0/24
EOF
        printf "${GREEN}✓ 已创建白名单文件: $WHITELIST_FILE${RESET}\n"
    fi
    
    # 初始化封禁列表
    if [ ! -f "$BANLIST_FILE" ]; then
        touch "$BANLIST_FILE"
        printf "${GREEN}✓ 已创建封禁列表文件: $BANLIST_FILE${RESET}\n"
    fi
    
    # 创建日志文件
    if [ ! -f "$LOG_FILE" ]; then
        touch "$LOG_FILE"
        printf "${GREEN}✓ 已创建日志文件: $LOG_FILE${RESET}\n"
    fi
}

# 智能检测SSH日志文件
detect_ssh_log() {
    # 预定义的日志路径（使用字符串而非数组以兼容sh）
    local possible_logs="/var/log/auth.log /var/log/secure /var/log/messages /var/log/syslog /var/log/journal/ssh.log /var/log/authlog"
    
    log_message "INFO" "正在智能检测SSH日志文件..."
    
    # 首先尝试配置文件中的路径
    if [ -n "$SSH_LOG" ] && [ -f "$SSH_LOG" ]; then
        log_message "INFO" "使用配置文件中的SSH日志: $SSH_LOG"
        return 0
    fi
    
    # 尝试预定义的路径
    for log_path in $possible_logs; do
        if [ -f "$log_path" ]; then
            # 检查文件是否包含SSH相关内容
            if grep -q -E "(sshd|SSH|Failed password|Invalid user)" "$log_path" 2>/dev/null; then
                SSH_LOG="$log_path"
                log_message "INFO" "检测到SSH日志文件: $SSH_LOG"
                return 0
            fi
        fi
    done
    
    # 使用find命令搜索可能的SSH日志
    log_message "INFO" "使用find命令搜索SSH日志文件..."
    local found_logs=$(find /var/log -name "*auth*" -o -name "*secure*" -o -name "*ssh*" 2>/dev/null)
    
    for log_path in $found_logs; do
        if [ -f "$log_path" ] && [ -r "$log_path" ]; then
            # 检查文件是否包含SSH相关内容
            if grep -q -E "(sshd|SSH|Failed password|Invalid user)" "$log_path" 2>/dev/null; then
                SSH_LOG="$log_path"
                log_message "INFO" "找到SSH日志文件: $SSH_LOG"
                return 0
            fi
        fi
    done
    
    # 检查systemd journal
    if command -v journalctl >/dev/null 2>&1; then
        if journalctl -u ssh 2>/dev/null | head -1 >/dev/null 2>&1; then
            SSH_LOG="journalctl"
            log_message "INFO" "使用systemd journal获取SSH日志"
            return 0
        elif journalctl -u sshd 2>/dev/null | head -1 >/dev/null 2>&1; then
            SSH_LOG="journalctl"
            log_message "INFO" "使用systemd journal获取SSH日志"
            return 0
        fi
    fi
    
    return 1
}

# 获取SSH日志内容（支持文件和journal）
get_ssh_log_content() {
    local hours_ago="$1"
    
    if [ "$SSH_LOG" = "journalctl" ]; then
        # 使用systemd journal
        journalctl -u ssh -u sshd --since "$hours_ago hours ago" 2>/dev/null || \
        journalctl --since "$hours_ago hours ago" | grep -E "(sshd|SSH)" 2>/dev/null
    else
        # 使用日志文件
        if [ -f "$SSH_LOG" ]; then
            cat "$SSH_LOG"
        else
            echo ""
        fi
    fi
}

# 加载配置
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        source "$CONFIG_FILE"
    else
        printf "${RED}❌ 配置文件不存在，请先运行初始化${RESET}\n"
        exit 1
    fi
    
    # 智能检测SSH日志文件
    if ! detect_ssh_log; then
        printf "${RED}❌ 未找到SSH日志文件${RESET}\n"
        printf "${YELLOW}支持的系统和日志路径:${RESET}\n"
        printf "  Ubuntu/Debian/Deepin: /var/log/auth.log\n"
        printf "  CentOS/RHEL/Rocky: /var/log/secure\n"
        printf "  通用: /var/log/messages\n"
        printf "  systemd: journalctl -u sshd\n\n"
        
        printf "${CYAN}诊断建议:${RESET}\n"
        printf "1. 检查SSH服务是否正在运行: ${YELLOW}systemctl status sshd${RESET}\n"
        printf "2. 检查系统日志配置: ${YELLOW}ls -la /var/log/${RESET}\n"
        printf "3. 手动查看SSH日志: ${YELLOW}journalctl -u sshd${RESET}\n"
        printf "4. 检查rsyslog配置: ${YELLOW}cat /etc/rsyslog.conf${RESET}\n\n"
        
        # 尝试提供更多帮助
        printf "${MAGENTA}如果是Deepin系统，可以尝试:${RESET}\n"
        printf "  ${YELLOW}sudo journalctl -u sshd | grep Failed${RESET}\n"
        printf "  ${YELLOW}find /var/log -name '*' | xargs grep -l 'sshd' 2>/dev/null${RESET}\n"
        
        exit 1
    fi
    
    # 检测防火墙工具
    if ! detect_firewall; then
        printf "${RED}❌ 防火墙工具检测失败${RESET}\n"
        exit 1
    fi
}

# 记录日志
log_message() {
    local level="$1"
    local message="$2"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] [$level] $message" >> "$LOG_FILE"
    
    case "$level" in
        "INFO")
            printf "${GREEN}[INFO]${RESET} $message\n"
            ;;
        "WARN")
            printf "${YELLOW}[WARN]${RESET} $message\n"
            ;;
        "ERROR")
            printf "${RED}[ERROR]${RESET} $message\n"
            ;;
        "BAN")
            printf "${RED}[BAN]${RESET} $message\n"
            ;;
        *)
            printf "[LOG] $message\n"
            ;;
    esac
}

# 检查IP是否在白名单中
is_whitelisted() {
    local ip="$1"
    
    # 检查基本白名单
    if grep -q "^$ip$" "$WHITELIST_FILE" 2>/dev/null; then
        return 0
    fi
    
    # 检查CIDR格式白名单
    while read -r line; do
        # 跳过注释和空行
        case "$line" in
            \#*) continue ;;
            "") continue ;;
        esac
        
        # 检查CIDR格式
        case "$line" in
            */[0-9]*)
                if command -v ipcalc >/dev/null 2>&1; then
                    if ipcalc -c "$line" 2>/dev/null && ipcalc -c "$ip" "$line" 2>/dev/null; then
                        return 0
                    fi
                fi
                ;;
        esac
    done < "$WHITELIST_FILE"
    
    return 1
}

# 设置自动解封
setup_auto_unban() {
    local ip="$1"
    local duration="$2"
    
    if [ -n "$duration" ] && [ "$duration" -gt 0 ]; then
        # 使用at命令设置定时解封
        if command -v at >/dev/null 2>&1; then
            # 创建临时解封脚本
            local unban_script="/tmp/unban_${ip}_$(date +%s).sh"
            cat > "$unban_script" << EOF
#!/bin/bash
# 自动解封脚本
cd "$SCRIPT_DIR"
source "$CONFIG_FILE"

# 设置防火墙工具变量
FIREWALL_TOOL="$FIREWALL_TOOL"

# 解封IP
case "\$FIREWALL_TOOL" in
    "iptables")
        while iptables -D INPUT -s $ip -j DROP 2>/dev/null; do
            continue
        done
        ;;
    "nftables")
        handles=\$(nft -a list chain inet filter input | grep "$ip" | awk '{print \$NF}')
        for handle in \$handles; do
            nft delete rule inet filter input handle "\$handle" 2>/dev/null
        done
        ;;
    "ufw")
        ufw delete deny from "$ip" 2>/dev/null
        ;;
esac

# 从封禁列表中移除
if [ -f "$BANLIST_FILE" ]; then
    grep -v "|$ip|" "$BANLIST_FILE" > "$BANLIST_FILE.tmp" && mv "$BANLIST_FILE.tmp" "$BANLIST_FILE"
fi

# 记录日志
echo "\$(date '+%Y-%m-%d %H:%M:%S') [INFO] 自动解封IP: $ip" >> "$LOG_FILE"

# 清理临时脚本
rm -f "$unban_script"
EOF
            chmod +x "$unban_script"
            
            # 使用at命令调度执行
            echo "$unban_script" | at now + $duration minutes 2>/dev/null
            log_message "INFO" "已设置 $duration 分钟后自动解封 $ip"
        else
            log_message "WARN" "at命令不可用，无法设置自动解封"
        fi
    fi
}

# 发送通知
send_notification() {
    local title="$1"
    local message="$2"
    
    # 邮件通知
    if [ "$ENABLE_EMAIL" = "true" ] && [ -n "$EMAIL_TO" ]; then
        if command -v mail >/dev/null 2>&1; then
            echo -e "$message" | mail -s "$title" "$EMAIL_TO"
        fi
    fi
    
    # 钉钉通知
    if [ "$ENABLE_DINGTALK" = "true" ] && [ -n "$DINGTALK_WEBHOOK" ]; then
        if command -v curl >/dev/null 2>&1; then
            curl -s -X POST "$DINGTALK_WEBHOOK" \
                -H "Content-Type: application/json" \
                -d "{\"msgtype\":\"text\",\"text\":{\"content\":\"$title\n$message\"}}" >/dev/null
        fi
    fi
}

# 分析SSH日志
analyze_ssh_log() {
    log_message "INFO" "开始分析SSH登录日志"
    
    # 检查日志源
    if [ "$SSH_LOG" = "journalctl" ]; then
        log_message "INFO" "使用systemd journal分析SSH日志"
    elif [ ! -f "$SSH_LOG" ]; then
        log_message "ERROR" "SSH日志文件不存在: $SSH_LOG"
        return 1
    else
        log_message "INFO" "分析SSH日志文件: $SSH_LOG"
    fi
    
    # 计算时间范围
    local check_time=$(date -d "$CHECK_PERIOD_HOURS hours ago" '+%b %d %H:%M:%S')
    local current_year=$(date '+%Y')
    
    # 获取SSH日志内容并分析
    local ssh_log_content
    if [ "$SSH_LOG" = "journalctl" ]; then
        ssh_log_content=$(get_ssh_log_content "$CHECK_PERIOD_HOURS")
    else
        ssh_log_content=$(cat "$SSH_LOG")
    fi
    
    # 创建临时awk脚本来避免shell解析问题
    cat > "/tmp/ssh_analysis.awk" << 'AWKEOF'
    # 处理不同格式的SSH失败日志
    /Failed password/ || /Invalid user/ || /authentication failure/ {
        # 提取时间戳
        if ($1 ~ /^[0-9]{4}-[0-9]{2}-[0-9]{2}/) {
            # systemd journal 格式: 2023-06-10 14:39:21
            time_str = $1 " " $2
        } else {
            # 传统 syslog 格式: Jun 10 14:39:21
            time_str = $1 " " $2 " " $3
        }
        
        # 提取IP地址
        ip = ""
        if (match($0, /from ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/, arr)) {
            ip = arr[1]
        }
        
        # 提取用户名
        user = ""
        if (match($0, /for ([a-zA-Z0-9_-]+)/, arr)) {
            user = arr[1]
        } else if (match($0, /user ([a-zA-Z0-9_-]+)/, arr)) {
            user = arr[1]
        }
        
        if (ip != "" && time_str >= time_limit) {
            count[ip]++
            users[ip] = users[ip] ? users[ip] "," user : user
            last_time[ip] = time_str
        }
    }
    END {
        for (ip in count) {
            if (count[ip] >= threshold) {
                print ip "|" count[ip] "|SSH暴力破解|" users[ip] "|" last_time[ip]
            }
        }
    }
AWKEOF
    
    # 检测SSH暴力破解
    echo "$ssh_log_content" | awk -v threshold="$MAX_FAILED_ATTEMPTS" -v time_limit="$check_time" -v year="$current_year" -f "/tmp/ssh_analysis.awk" > "/tmp/ssh_violations.txt"
    
    # 检测敏感用户登录尝试
    if [ -n "$SENSITIVE_USERS" ]; then
        for sensitive_user in $(echo "$SENSITIVE_USERS" | tr ',' ' '); do
            # 创建敏感用户检测的awk脚本
            cat > "/tmp/sensitive_user.awk" << 'AWKEOF'
            (/Failed password/ || /Invalid user/) && $0 ~ user {
                if (match($0, /from ([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})/, arr)) {
                    ip = arr[1]
                    if ($1 ~ /^[0-9]{4}-[0-9]{2}-[0-9]{2}/) {
                        time_str = $1 " " $2
                    } else {
                        time_str = $1 " " $2 " " $3
                    }
                    if (time_str >= time_limit) {
                        count[ip]++
                    }
                }
            }
            END {
                for (ip in count) {
                    if (count[ip] >= 2) {
                        print ip "|" count[ip] "|敏感用户登录尝试:" user
                    }
                }
            }
AWKEOF
            echo "$ssh_log_content" | awk -v user="$sensitive_user" -v time_limit="$check_time" -f "/tmp/sensitive_user.awk" >> "/tmp/ssh_violations.txt"
        done
    fi
    
    # 处理违规IP
    if [ -f "/tmp/ssh_violations.txt" ]; then
        while IFS='|' read -r ip count reason user_list last_time; do
            if [ -n "$ip" ]; then
                ban_ip "$ip" "$reason" "$count"
            fi
        done < "/tmp/ssh_violations.txt"
        rm -f "/tmp/ssh_violations.txt"
    fi
    
    # 清理临时文件
    rm -f "/tmp/ssh_analysis.awk" "/tmp/sensitive_user.awk"
    
    log_message "INFO" "SSH日志分析完成"
}

# 显示主菜单
show_menu() {
    printf "\n${BLUE}========== SSH监控与封禁系统 ==========${RESET}\n"
    printf "${GREEN}1) 📊 查看当前状态${RESET}\n"
    printf "${GREEN}2) 🔍 手动检测SSH攻击${RESET}\n"
    printf "${GREEN}3) 🚫 手动封禁IP${RESET}\n"
    printf "${GREEN}4) ✅ 解封IP${RESET}\n"
    printf "${GREEN}5) 📝 管理白名单${RESET}\n"
    printf "${GREEN}6) 📋 查看封禁列表${RESET}\n"
    printf "${GREEN}7) 📈 查看日志${RESET}\n"
    printf "${GREEN}8) ⚙️  配置设置${RESET}\n"
    printf "${GREEN}9) 🔄 安装定时任务${RESET}\n"
    printf "${GREEN}10) 📊 生成安全报告${RESET}\n"
    printf "${YELLOW}0) 🚪 退出${RESET}\n"
    printf "${BLUE}=================================${RESET}\n"
    read -p "请选择操作 (0-10): " choice
}

# 查看当前状态
show_status() {
    printf "\n${CYAN}========== SSH安全状态 ==========${RESET}\n"
    
    # 当前封禁的IP数量
    local banned_count=$(get_banned_count)
    printf "${GREEN}当前封禁IP数量: ${banned_count}${RESET}\n"
    
    # 白名单IP数量
    local whitelist_count=$(grep -v '^#' "$WHITELIST_FILE" | grep -v '^$' | wc -l)
    printf "${GREEN}白名单IP数量: ${whitelist_count}${RESET}\n"
    
    # 今日封禁记录
    local today=$(date '+%Y-%m-%d')
    local today_bans=$(grep "$today" "$BANLIST_FILE" 2>/dev/null | wc -l)
    printf "${GREEN}今日封禁次数: ${today_bans}${RESET}\n"
    
    # SSH服务状态
    if systemctl is-active ssh >/dev/null 2>&1 || systemctl is-active sshd >/dev/null 2>&1; then
        printf "${GREEN}SSH服务状态: ${GREEN}运行中${RESET}\n"
    else
        printf "${GREEN}SSH服务状态: ${RED}未运行${RESET}\n"
    fi
    
    # 最近SSH失败尝试
    printf "\n${YELLOW}最近SSH失败尝试 (最多10条):${RESET}\n"
    if [ -f "$SSH_LOG" ]; then
        tail -100 "$SSH_LOG" | grep -E "(Failed password|Invalid user)" | tail -10 | while read line; do
            local ip=$(echo "$line" | grep -oE '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | head -1)
            local time=$(echo "$line" | awk '{print $1, $2, $3}')
            printf "  ${RED}$ip${RESET} - $time\n"
        done
    fi
    
    # 最近封禁的IP
    printf "\n${YELLOW}最近封禁的IP (最多10个):${RESET}\n"
    tail -10 "$BANLIST_FILE" 2>/dev/null | while IFS='|' read -r timestamp ip reason count; do
        printf "  ${RED}$ip${RESET} - $timestamp - $reason\n"
    done
    
    read -p "按Enter键返回主菜单..." dummy
}

# 手动检测SSH攻击
manual_scan() {
    printf "\n${CYAN}========== 手动检测SSH攻击 ==========${RESET}\n"
    printf "${YELLOW}正在分析SSH日志文件...${RESET}\n"
    
    analyze_ssh_log
    
    printf "${GREEN}✓ 检测完成${RESET}\n"
    read -p "按Enter键返回主菜单..." dummy
}

# 手动封禁IP
manual_ban() {
    printf "\n${CYAN}========== 手动封禁IP ==========${RESET}\n"
    read -p "请输入要封禁的IP地址: " ip
    read -p "请输入封禁原因: " reason
    
    if [ -n "$ip" ] && [ -n "$reason" ]; then
        ban_ip "$ip" "手动封禁: $reason" "手动"
    else
        printf "${RED}❌ IP地址和原因不能为空${RESET}\n"
    fi
    
    read -p "按Enter键返回主菜单..." dummy
}

# 解封IP
manual_unban() {
    printf "\n${CYAN}========== 解封IP ==========${RESET}\n"
    
    # 显示当前封禁的IP
    printf "${YELLOW}当前封禁的IP列表:${RESET}\n"
    show_current_banned_ips
    
    printf "\n"
    read -p "请输入要解封的IP地址: " ip
    
    if [ -n "$ip" ]; then
        unban_ip "$ip"
        printf "${GREEN}✓ 操作完成${RESET}\n"
    else
        printf "${RED}❌ IP地址不能为空${RESET}\n"
    fi
    
    read -p "按Enter键返回主菜单..." dummy
}

# 管理白名单
manage_whitelist() {
    while true; do
        printf "\n${CYAN}========== 白名单管理 ==========${RESET}\n"
        printf "${GREEN}1) 查看白名单${RESET}\n"
        printf "${GREEN}2) 添加IP到白名单${RESET}\n"
        printf "${GREEN}3) 从白名单移除IP${RESET}\n"
        printf "${GREEN}4) 编辑白名单文件${RESET}\n"
        printf "${YELLOW}0) 返回主菜单${RESET}\n"
        read -p "请选择操作: " choice
        
        case $choice in
            1)
                printf "\n${YELLOW}当前白名单:${RESET}\n"
                cat "$WHITELIST_FILE"
                read -p "按Enter键继续..." dummy
                ;;
            2)
                read -p "请输入要添加的IP地址: " ip
                if [ -n "$ip" ]; then
                    echo "$ip" >> "$WHITELIST_FILE"
                    printf "${GREEN}✓ 已添加到白名单${RESET}\n"
                fi
                read -p "按Enter键继续..." dummy
                ;;
            3)
                read -p "请输入要移除的IP地址: " ip
                if [ -n "$ip" ]; then
                    grep -v "^$ip$" "$WHITELIST_FILE" > "$WHITELIST_FILE.tmp"
                    mv "$WHITELIST_FILE.tmp" "$WHITELIST_FILE"
                    printf "${GREEN}✓ 已从白名单移除${RESET}\n"
                fi
                read -p "按Enter键继续..." dummy
                ;;
            4)
                if command -v nano >/dev/null 2>&1; then
                    nano "$WHITELIST_FILE"
                elif command -v vi >/dev/null 2>&1; then
                    vi "$WHITELIST_FILE"
                else
                    printf "${RED}❌ 未找到可用的编辑器${RESET}\n"
                fi
                ;;
            0)
                break
                ;;
            *)
                printf "${RED}❌ 无效选项${RESET}\n"
                ;;
        esac
    done
}

# 查看封禁列表
show_banlist() {
    printf "\n${CYAN}========== SSH攻击封禁列表 ==========${RESET}\n"
    
    if [ -s "$BANLIST_FILE" ]; then
        printf "${YELLOW}格式: 时间|IP地址|封禁原因|详细信息${RESET}\n"
        printf "${YELLOW}===============================================${RESET}\n"
        cat "$BANLIST_FILE" | tail -20
        
        local total_bans=$(wc -l < "$BANLIST_FILE")
        printf "\n${GREEN}总计封禁记录: ${total_bans} 条 (显示最近20条)${RESET}\n"
    else
        printf "${YELLOW}暂无封禁记录${RESET}\n"
    fi
    
    read -p "按Enter键返回主菜单..." dummy
}

# 查看日志
show_logs() {
    printf "\n${CYAN}========== SSH监控日志 ==========${RESET}\n"
    
    if [ -s "$LOG_FILE" ]; then
        tail -50 "$LOG_FILE"
        printf "\n${GREEN}显示最近50条日志记录${RESET}\n"
    else
        printf "${YELLOW}暂无日志记录${RESET}\n"
    fi
    
    read -p "按Enter键返回主菜单..." dummy
}

# 配置设置
configure_settings() {
    printf "\n${CYAN}========== 配置设置 ==========${RESET}\n"
    printf "${YELLOW}当前配置文件: $CONFIG_FILE${RESET}\n"
    printf "${YELLOW}1) 编辑配置文件${RESET}\n"
    printf "${YELLOW}2) 重载配置${RESET}\n"
    printf "${YELLOW}3) 测试SSH日志路径${RESET}\n"
    printf "${YELLOW}0) 返回${RESET}\n"
    read -p "请选择操作: " choice
    
    case $choice in
        1)
            if command -v nano >/dev/null 2>&1; then
                nano "$CONFIG_FILE"
            elif command -v vi >/dev/null 2>&1; then
                vi "$CONFIG_FILE"
            else
                printf "${RED}❌ 未找到可用的编辑器${RESET}\n"
            fi
            ;;
        2)
            load_config
            printf "${GREEN}✓ 配置已重载${RESET}\n"
            ;;
        3)
            printf "${YELLOW}当前SSH日志路径: $SSH_LOG${RESET}\n"
            if [ -f "$SSH_LOG" ]; then
                printf "${GREEN}✓ 日志文件存在${RESET}\n"
                printf "${YELLOW}最近5条SSH失败记录:${RESET}\n"
                tail -100 "$SSH_LOG" | grep -E "(Failed password|Invalid user)" | tail -5
            else
                printf "${RED}❌ 日志文件不存在${RESET}\n"
            fi
            ;;
        0)
            return
            ;;
        *)
            printf "${RED}❌ 无效选项${RESET}\n"
            ;;
    esac
    
    read -p "按Enter键返回主菜单..." dummy
}

# 安装定时任务
install_cron() {
    printf "\n${CYAN}========== 安装定时任务 ==========${RESET}\n"
    
    local cron_script="/usr/local/bin/ssh-monitor-cron.sh"
    
    # 创建定时任务脚本
    cat > "$cron_script" << EOF
#!/bin/bash
# SSH监控定时任务脚本
cd "$SCRIPT_DIR"

# 检查配置文件
if [ ! -f "$CONFIG_FILE" ]; then
    echo "\$(date '+%Y-%m-%d %H:%M:%S') [ERROR] 配置文件不存在" >> "$LOG_FILE"
    exit 1
fi

# 加载配置
source "$CONFIG_FILE"

# 检查iptables是否可用
if ! command -v iptables >/dev/null 2>&1; then
    echo "\$(date '+%Y-%m-%d %H:%M:%S') [ERROR] iptables命令不可用" >> "$LOG_FILE"
    exit 1
fi

# 执行监控检测
$0 --scan >/dev/null 2>&1
EOF
    
    chmod +x "$cron_script"
    
    # 添加到crontab
    (crontab -l 2>/dev/null | grep -v "ssh-monitor-cron.sh"; echo "# SSH监控定时任务 - 每3分钟执行一次"; echo "*/3 * * * * $cron_script") | crontab -
    
    printf "${GREEN}✓ 定时任务安装成功${RESET}\n"
    printf "${YELLOW}定时任务将每3分钟执行一次SSH攻击检测${RESET}\n"
    printf "${YELLOW}脚本路径: $cron_script${RESET}\n"
    
    read -p "按Enter键返回主菜单..." dummy
}

# 生成安全报告
generate_report() {
    local report_file="$SCRIPT_DIR/ssh_security_report_$(date +%Y%m%d).txt"
    
    printf "========== SSH安全监控日报 ==========\n" > "$report_file"
    printf "生成时间: $(date '+%Y-%m-%d %H:%M:%S')\n\n" >> "$report_file"
    
    # 今日统计
    local today=$(date '+%Y-%m-%d')
    local today_bans=$(grep "$today" "$BANLIST_FILE" 2>/dev/null | wc -l)
    local current_banned=$(get_banned_count)
    
    printf "今日SSH攻击封禁次数: %d\n" "$today_bans" >> "$report_file"
    printf "当前封禁IP数量: %d\n\n" "$current_banned" >> "$report_file"
    
    # 攻击类型统计
    printf "攻击类型统计:\n" >> "$report_file"
    grep "$today" "$BANLIST_FILE" 2>/dev/null | cut -d'|' -f3 | sort | uniq -c | sort -nr >> "$report_file"
    printf "\n" >> "$report_file"
    
    # Top 10 攻击IP
    printf "Top 10 SSH攻击IP:\n" >> "$report_file"
    grep "$today" "$BANLIST_FILE" 2>/dev/null | cut -d'|' -f2 | sort | uniq -c | sort -nr | head -10 >> "$report_file"
    printf "\n" >> "$report_file"
    
    # 攻击时间分布
    printf "攻击时间分布:\n" >> "$report_file"
    grep "$today" "$BANLIST_FILE" 2>/dev/null | cut -d'|' -f1 | cut -d' ' -f2 | cut -d':' -f1 | sort | uniq -c >> "$report_file"
    printf "\n" >> "$report_file"
    
    # 最近SSH失败尝试统计
    if [ -f "$SSH_LOG" ]; then
        printf "最近SSH失败尝试统计:\n" >> "$report_file"
        tail -1000 "$SSH_LOG" | grep "$today" | grep -E "(Failed password|Invalid user)" | \
        grep -oE '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}' | sort | uniq -c | sort -nr | head -15 >> "$report_file"
    fi
    
    log_message "INFO" "SSH安全报告已生成: $report_file"
    printf "${GREEN}✓ SSH安全报告已生成: $report_file${RESET}\n"
    
    read -p "按Enter键返回主菜单..." dummy
}

# 主程序入口
main() {
    # 检查参数
    case "$1" in
        --init)
            printf "${GREEN}正在初始化SSH监控系统...${RESET}\n"
            init_config
            load_config
            printf "${GREEN}✓ 初始化完成${RESET}\n"
            exit 0
            ;;
        --scan)
            load_config
            analyze_ssh_log
            exit 0
            ;;
        --help|-h)
            printf "${CYAN}SSH访问监控与自动封禁脚本${RESET}\n"
            printf "用法: $0 [选项]\n"
            printf "选项:\n"
            printf "  --init      初始化配置文件\n"
            printf "  --scan      扫描SSH日志并执行封禁\n"
            printf "  --help      显示此帮助信息\n"
            exit 0
            ;;
    esac
    
    # 检查配置文件
    if [ ! -f "$CONFIG_FILE" ]; then
        printf "${YELLOW}首次运行，正在初始化...${RESET}\n"
        init_config
    fi
    
    load_config
    
    # 显示欢迎信息
    printf "${GREEN}欢迎使用SSH访问监控与自动封禁系统！${RESET}\n"
    printf "${CYAN}配置目录: $SCRIPT_DIR${RESET}\n"
    printf "${CYAN}SSH日志: $SSH_LOG${RESET}\n"
    
    # 主循环
    while true; do
        show_menu
        
        case $choice in
            1) show_status ;;
            2) manual_scan ;;
            3) manual_ban ;;
            4) manual_unban ;;
            5) manage_whitelist ;;
            6) show_banlist ;;
            7) show_logs ;;
            8) configure_settings ;;
            9) install_cron ;;
            10) generate_report ;;
            0) 
                printf "${GREEN}感谢使用SSH监控系统！${RESET}\n"
                break 
                ;;
            *)
                printf "${RED}❌ 无效选项，请重新选择${RESET}\n"
                ;;
        esac
    done
}

# 运行主程序
main "$@" 